from ultralytics import YOLO

# Load a pretrained YOLO11n model
model = YOLO("runs/detect/train39/weights/last.pt")

# Train the model
train_results = model.train(
    data=r"datasets/traffic_demo1/traffic_demo1.yaml",  # Path to dataset configuration file
    epochs=10,  # Number of training epochs
    imgsz=320,  # Image size for training
    device="cpu",
    batch=8,  # Device to run on (e.g., 'cpu', 0, [0,1,2,3])
    workers=0,
)

# Evaluate the model's performance on the validation set
metrics = model.train()


# Export the model to ONNX format for deployment
path = model.export(format="onnx")  # Returns the path to the exported model